#!/usr/bin/env python3
# $Id: test3-7.py b3a8aba613bc 2013/04/26 05:38:41 pgurumur $
# $Author: pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

__version__  = "$Revision: b3a8aba613bc $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2013-04-26 05:38:41Z $"

import optparse
import os
import select
import signal
import socket
import struct
import sys
import time

import dns.resolver

from venya3.lib import cmdparser
from venya3.lib import exception
from venya3.lib import layer3
from venya3.lib import ip
from venya3.lib import stdlib

sig2name = dict((getattr(signal, num), num) \
          for num in dir(signal) if num.startswith('SIG') and '_' not in num)

class traceroute(cmdparser.GetOpt):
   option_list = (
         optparse.make_option("-4", "--ipv4", action = "store_true",
            help = "Traceroute to IPv4 address", default = False, dest = "v4"),
         optparse.make_option("-6", "--ipv6", action = "store_true",
            help = "Traceroute to IPv6 address", default = False, dest = "v6"),
         optparse.make_option("-n", "--nodns", action = "store_true",
            help = "No reverse lookup IPv4/IPv6 address in DNS",
            default = False, dest = "dns"),
   ) + cmdparser.GetOpt.option_list

   help = """

   This program will do traceroute of both IPv4 and IPv6 addresses/hostnames
   using UDP socket and relying on time-exceeded ICMP/ICMP6 return messages 
   and port unreachables from the actual hosts to get the path of the
   IPv4/IPv6 datagrams.

   In addition to the existing functionality that is in traditional
   traceroute programs available with the operating systems, this traceroute
   gives the AS number associated with the intermediary and final IP addresses
   and the company/entity names the AS numbers belong to with the help of
   cymru.com

   Examples:
   1. traceroute.py -4 www.google.com
   The above example would trace the path to Google's web server and prints
   out the intermediary and final AS numbers and its associated entries

   2. traceroute.py -4n www.google.com
   Same as example 1, except the program does not resolve the reverse records 
   for the IP addresses in the network path

   3. traceroute.py -6 www.google.com
   Same as example 1, except it works for IPv6

   4. traceroute.py -6n www.google.com
   Same as example 2, except it works for IPv6

   Example output:
   traceroute4 to www.google.com(74.125.224.145), pid: 16569, maxhops: 30
   01 172.29.17.1 (chrysalis.kashyapa.net) 0.6187 secs
   02 67.169.140.1 (c-67-169-140-1.hsd1.ca.comcast.net, ASN: 33651) 1.1694 secs
   03 68.85.190.13 (ge-3-8-ur01.fremont.ca.sfba.comcast.net, ASN: 33651) 1.7159 secs
   04 68.85.154.114 (te-1-8-0-3-ar01.sfsutro.ca.sfba.comcast.net, ASN: 33651) 2.2519 secs
   05 68.86.90.153 (he-1-7-0-0-cr01.sanjose.ca.ibone.comcast.net, ASN: 7922) 2.7977 secs
   06 68.86.87.142 (pos-0-3-0-0-pe01.529bryant.ca.ibone.comcast.net, ASN: 7922) 3.4216 secs
   07 66.208.228.226 (ASN: 7922) 3.9924 secs
   08 72.14.232.136 (ASN: 15169) 4.6567 secs
   09 64.233.174.119 (ASN: 15169) 5.2250 secs
   10 74.125.224.145 (nuq04s09-in-f17.1e100.net, ASN: 15169) 5.8542 secs
   ---------------------------------------
   ASN 33651, Entity: CMCS - Comcast Cable Communications, Inc.
   ASN 7922, Entity: COMCAST-7922 - Comcast Cable Communications, Inc.
   ASN 15169, Entity: GOOGLE - Google Inc.
   """

   def __init__(self, **Keywords):
      super(stdlib.my_class(self), self).__init__(**Keywords)
      self._port = 33434
      self._maxhops = 30
      self._ttl = 1
      self._rsock = None
      self._usock = None
      self._family = None
      self._dns = True
      self._debug = False
      self._asList = []
      self._icmptype = 0
      self._icmpcode = 0
      self._origin = {}
      self._asn = {}
      self._seconds = 0

   def debug(self, fmt, *args):
      if self._debug:
         print(stdlib.sprintf(fmt, *args))

   def sighandle(self, signum, frame):
      print("\n")
      if self._rsock:
         self._rsock.close()

      if self._usock:
         self._usock.close()

      self.error("program with pid %d terminated with signal: %s(%d)",
            os.getpid(), sig2name[signum], signum)

   def setsignal(self):
      try:
         signal.signal(signal.SIGTERM, self.sighandle)
         signal.signal(signal.SIGINT, self.sighandle)
         signal.signal(signal.SIGHUP, self.sighandle)
      except ValueError as err:
         pass

   def dnsquery(self, record, rtype = "TXT"):
      data = None
      retval = None
      try:
         data = dns.resolver.query(record, rtype)
      except dns.resolver.NXDOMAIN:
         pass
      else:
         retval = str(data.rrset).split("TXT")[1].replace("\"", "").split("|")

      return retval

   def contains(self, address):
      retval = 0
      try:
         for ix in self._origin.keys():
            if address in ip.ip(ix):
               retval = self._origin[ix]
               break
      except exception.IPError as err:
         pass

      return retval

   def getasn(self, address):
      retval = 0
      asn = self.contains(address)
      temp = None
      try:
         temp = ip.ip(address)
      except exception.IPError as err:
         pass

      if temp and not temp.private:
         if asn == 0:
            if self.version == 4:
               query = "%s.origin.asn.cymru.com" %stdlib.ip2ptr(address)
            elif self.version == 6:
               query = "%s.origin6.asn.cymru.com" %stdlib.ip2ptr6(address)

            data = self.dnsquery(query)
            if data:
               retval = int(data[0].strip())
               self._origin[data[1].strip()] = retval
               if retval not in self._asList:
                  self._asList.append(retval)
         else:
            retval = asn

      return retval

   def getentity(self, asn):
      retval = None
      if asn in self._asn:
         retval = self._asn[asn]
      else:
         query = "AS%d.asn.cymru.com" %asn
         data = self.dnsquery(query)
         if data:
            retval = data[-1].strip()
            self._asn[asn] = retval

      return retval

   def _icmpsocket(self):
      sockList = []
      try:
         protocol = None
         if self.version == 4:
            protocol = socket.getprotobyname('icmp')
         elif self.version == 6:
            protocol = socket.getprotobyname('ipv6-icmp')

         self.debug("protocol is %d", protocol)
         self._rsock = socket.socket(self._family, socket.SOCK_RAW, protocol)
         if self.version == 4:
            self._rsock.setsockopt(socket.IPPROTO_IP, socket.IP_HDRINCL, 1)
         elif self.version == 6:
            self._rsock.setsockopt(socket.IPPROTO_IPV6,
                  socket.IPV6_RECVPKTINFO, 1)
            self._rsock.setsockopt(socket.IPPROTO_IPV6,
                  socket.IPV6_RECVHOPLIMIT, 1)
      except socket.error as err:
         self.error("error creating raw sockets: %s", err.args[0])
      else:
         self.debug("icmp socket fileno: %d", self._rsock.fileno())

   def _udpsocket(self, host):
      sockList = []
      if self._rsock:
         if self.version == 4:
            self._rsock.bind(("", 0))
         elif self.version == 6:
            self._rsock.bind(("", 0))

      try:
         sockList = socket.getaddrinfo(host, 0, self._family,
               socket.SOCK_DGRAM, socket.IPPROTO_UDP)
      except socket.gaierror as err:
         self.error("getaddrinfo: errno %d, message: %s", err.args[0],
               err.args[1])

      mesg = None
      for res in sockList:
         af, socktype, proto, canonname, sa = res
         try:
            self._usock = socket.socket(af, socktype, proto)
         except socket.error:
            self._usock = None

         self.debug("udp socket number: %d", self._usock.fileno())
         try:
            if self.version == 4:
               self._usock.setsockopt(socket.SOL_IP, socket.IP_TTL, self._ttl)
            elif self.version == 6:
               self._usock.setsockopt(socket.IPPROTO_IPV6,
                     socket.IPV6_UNICAST_HOPS, self._ttl)

            self._usock.sendto("".encode("ascii"), (host, self._port))
         except socket.error as err:
            mesg = err.args[1]
            if self._usock:
               self._usock.close()

            self._usock = None
            continue

      if mesg:
         self.error("Cannot create UDP socket: %s", mesg)

   def _checkudp(self, data):
      retval = False
      test = None
      try:
         if len(data) == 56:
            test = struct.unpack(">HHHH", data[48:56])
      except struct.error as err:
         pass
      else:
         if test and len(test) and int(test[1]) == self._port:
            if (self.version == 4 and self._icmptype == 3 and
                  self._icmpcode == 3):
               retval = True
            elif (self.version == 6 and self._icmptype == 1 and
                  self._icmpcode == 4):
               retval = True

      return retval

   def _readicmp(self):
      retval = None
      try:
         rfd, wfd, efd = select.select([self._rsock], [], [], 1)
      except select.error as err:
         pass

      if self._rsock in rfd:
         curraddr = None
         data = None
         try:
            data, curraddr = self._rsock.recvfrom(500)
            curraddr = curraddr[0]
            try:
               header = None
               if self.version == 4:
                  header = data[20:28]
                  itype, code, chksum, temp = struct.unpack(">BBHI", header)
               elif self.version == 6:
                  header = data[0:4]
                  itype, code, chksum = struct.unpack("!BBH", header)
            except struct.error as err:
               pass
            else:
               self._icmptype = itype
               self._icmpcode = code
         except socket.error:
            pass
         else:
            if self.version == 4:
               if self._icmptype == 11 or (self._icmptype == 3 and
                     self._checkudp(data)):
                  retval = curraddr
               else:
                  retval = self._readicmp()
            elif self.version == 6:
               if self._icmptype == 3 or (self._icmptype == 1 and 
                     self._checkudp(data)):
                  retval = curraddr
               else:
                  retval = self._readicmp()

      return retval

   @property
   def version(self):
      retval = 0
      if self._family == socket.AF_INET:
         retval = 4
      elif self._family == socket.AF_INET6:
         retval = 6

      return retval

   def _reset(self):
      self._ttl = 1
      self._seconds = 0
      self._icmptype = 0
      self._icmpcode = 0
      self._asList = []

   def _tracepath(self, host):
      if not self._family:
         answers = None
         for rtype in ['A', 'AAAA']:
            try:
               answers = dns.resolver.query(host, rtype)
            except dns.resolver.NoAnswer:
               pass
            else:
               for rdata in answers:
                  atype = int(rdata.rdtype)
                  if atype == 1:
                     self._family = socket.AF_INET
                  elif atype == 28:
                     self._family = socket.AF_INET6

               if self._family:
                  break

      if host and self._family:
         mesg = stdlib.sprintf("\ntraceroute%d to %s(%s), pid: %d, maxhops: %d",
               self.version, host, layer3.host2ip(host, self.version),
               os.getpid(), self._maxhops)
         print(mesg)
         lastip = None 
         curraddr = None

         flag = True
         while flag:
            ix = time.time()
            self._icmpsocket()
            self._udpsocket(host)
            lastip = curraddr
            curraddr = self._readicmp()
            self._rsock.close()
            self._usock.close()

            currhost = None
            self._seconds += time.time() - ix
            if curraddr:
               currname = None
               asn = self.getasn(curraddr)
               if self._dns:
                  try:
                     currname = socket.gethostbyaddr(curraddr)[0]
                  except socket.error:
                     pass

               if asn:
                  if currname:
                     currhost = "%s (%s, ASN: %d)" %(curraddr, currname,
                           asn)
                  else:
                     currhost = "%s (ASN: %d)" %(curraddr, asn)
               else:
                  if currname:
                     currhost = "%s (%s)" %(curraddr, currname)
                  else:
                     currhost = curraddr
            else:
               currhost = "*"

            print("%02d %s %0.04f secs" %(self._ttl, currhost, self._seconds))
            self._ttl += 1

            if (self.version == 4 and self._icmptype == 3
               and self._icmpcode == 3):
                  flag = False
            elif (self.version == 6 and self._icmptype == 1
               and self._icmpcode == 4):
                  flag = False
            elif curraddr == layer3.host2ip(host,
                  self.version) or self._ttl > self._maxhops:
               flag = False

            if flag == False:
               print("---------------------------------------")
               for ix in self._asList:
                  entity = self.getentity(ix)
                  if entity:
                     print("ASN %d, Entity: %s" %(ix, entity))

   @cmdparser.getopts
   def handle(self):
      if self.options.v4 and self.options.v6:
         self.error("choose either IPv4 or IPv6, not both")

      try:
         host = self.arguments[1:]
      except IndexError:
         self.error("Need hostnames to traceroute!")

      if os.geteuid() != 0:
         self.error("Need to run as root!")

      self.setsignal()
      if self.options.dns:
         self._dns = False

      if self.options.v4:
         self._family = socket.AF_INET
      elif self.options.v6:
         self._family = socket.AF_INET6

      for ix in host:
         self._tracepath(ix)
         self._reset()

if __name__ == "__main__":
   client = None
   try:
      client = traceroute(version = 0.2, args = sys.argv, usage = "[ hosts ]")
      if client.length <= 1:
         client.printhelp()
      else:
         client.handle()
   except exception.GetOptError as err:
      print(err)
      sys.exit(1)
